/**
 * Copyright 2021 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "composite/optimize/axis_attr_normalize.h"

namespace akg {
class AxisAttrNormalizer : public IRMutator {
 public:
  Stmt Mutate_(const AttrStmt *op, const Stmt &s) {
    if (op->attr_key == "attrs") {
      auto attrs = Downcast<Map<std::string, NodeRef>>(op->node);
      if (attrs.find("axis") != attrs.end()) {
        Stmt body = IRMutator::Mutate(op->body);
        if (attrs["axis"].as<IntImm>()) {
          int64_t idx = axis_len_ > 0 ? (attrs["axis"].as<IntImm>()->value + axis_len_) % axis_len_ : 0;
          Array<Expr> new_axis = {make_const(Int(32), idx)};
          attrs.Set("axis", new_axis);
        } else {
          Array<Expr> axis = Downcast<Array<Expr>>(attrs["axis"]);
          Array<Expr> new_axis;
          for (auto val : axis) {
            auto imm = val.as<IntImm>();
            CHECK(imm);
            if (imm->value >= 0) {
              new_axis.push_back(val);
            } else {
              new_axis.push_back(make_const(Int(32), imm->value + axis_len_));
            }
            attrs.Set("axis", new_axis);
          }
        }
        return AttrStmt::make(attrs, op->attr_key, op->value, body);
      }
    }
    return s;
  }

  Stmt Mutate_(const Provide *op, const Stmt &s) {
    auto prim = op->value.as<Call>();
    CHECK(prim && prim->args.size() > 0);
    auto input = prim->args[0].as<Call>();
    if (input) {
      axis_len_ = input->args.size();
    }
    return s;
  }

 private:
  int64_t axis_len_{0};
};

Stmt AxisAttrNormalize::Run(const Stmt &stmt) { return AxisAttrNormalizer().Mutate(stmt); }
}  // namespace akg
